# 导入库（按功能分类）
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sqlalchemy import create_engine
import statsmodels.api as sm

class FinancialAnalyzer:
    def __init__(self, db_config):
        self.db_config = db_config
        self.engine = create_engine(
            f"mysql+pymysql://{db_config['user']}:{db_config['password']}"
            f"@{db_config['host']}:{db_config['port']}/{db_config['database']}"
            f"?charset={db_config['charset']}"
        )
        
    def _load_data(self):
        """加载并合并数据集"""
        query = """
        SELECT d.*, m.net_mf_vol, m.sell_elg_vol, m.buy_elg_vol, 
               m.sell_lg_vol, m.buy_lg_vol, i.closes as i_closes, i.vol as i_vol
        FROM date_1 d
        JOIN moneyflows m ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
        LEFT JOIN index_daily i ON d.trade_date = i.trade_date AND i.ts_code = '399001.SZ'
        WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' 
          AND d.ts_code = '000001.SZ'
        """
        return pd.read_sql(query, self.engine)

    def _calculate_returns(self, df):
        """计算收益率指标"""
        # 保留原始计算逻辑
        df['zd_close'] = df['closes'].shift(1)  # 注意此处疑似原始代码逻辑问题
        df['zs_closes'] = (df['i_closes'] - df['i_closes'].shift(1)) / df['i_closes'].shift(1)).round(2)
        return df.dropna(subset=['zd_close', 'zs_closes']).reset_index(drop=True)

    def _perform_pca(self, features):
        """执行主成分分析"""
        cov_matrix = np.cov(features, rowvar=False)
        eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)
        return eigenvalues, eigenvectors

    def _build_regression_model(self, X, y):
        """构建并训练回归模型"""
        X = sm.add_constant(X)
        model = sm.OLS(y, X)
        return model.fit()

    def _generate_plots(self, X, y, eigenvectors):
        """生成可视化图表"""
        # 散点图矩阵
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        for i, col in enumerate(X.columns[:3]):
            axes[i].scatter(X[col], y, s=50, alpha=0.7)
            axes[i].set_xlabel(col)
            axes[i].set_ylabel('(y)')
            axes[i].set_title(f'{col} vs Y')
        plt.tight_layout()
        plt.show()

        # 主成分方程输出
        for k in range(5):
            if k in [2, 4]: continue
            equation = f'CP{k+1} = '
            coeffs = [f'{round(val,2)}*X_{i+1}' for i, val in enumerate(eigenvectors[k])]
            equation += ' + '.join(coeffs).replace('+ -', '- ')
            print(equation)

    def analyze(self):
        """执行完整分析流程"""
        # 数据准备
        df = self._load_data()
        processed_df = self._calculate_returns(df)
        
        # 特征选择
        features = processed_df[['buy_lg_vol', 'sell_lg_vol', 
                                'sell_elg_vol', 'net_mf_vol', 'zs_closes']]
        
        # 主成分分析
        eigenvalues, eigenvectors = self._perform_pca(features)
        print(f'累计贡献率: {round(eigenvalues[:5].sum()/eigenvalues.sum(),4)*100}%')
        
        # 主成分计算
        principal_components = np.dot(features, eigenvectors[:, :5])
        pca_df = pd.concat([processed_df, pd.DataFrame(principal_components, 
                          columns=[f'PC{i+1}' for i in range(5)])], axis=1)

        # 全主成分回归
        full_model = self._build_regression_model(pca_df[['PC1','PC2','PC3','PC4','PC5']], 
                                                processed_df['zd_close'])
        print("\n全主成分回归结果:")
        print(full_model.summary())

        # 部分主成分回归
        selected_model = self._build_regression_model(pca_df[['PC1','PC2','PC3']], 
                                                    processed_df['zd_close'])
        print("\n选择主成分回归结果:")
        print(selected_model.summary())

        # 可视化
        self._generate_plots(pca_df[['PC1','PC2','PC3']], 
                            processed_df['zd_close'], 
                            eigenvectors)

# 配置参数
DB_CONFIG = {
    'host': '127.0.0.1',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tushare1',
    'port': 3306,
    'charset': 'utf8mb4'
}

if __name__ == "__main__":
    analyzer = FinancialAnalyzer(DB_CONFIG)
    analyzer.analyze()